#include <gtest/gtest.h>
#include <string>
#include <iostream>

#include "model/root.h"

const std::string deepseek_vl2_root = "../3rd/deepseek-vl2-tiny";

using namespace iangels;

TEST(ModelRootTest, InitTest_findFile_SUCCESS) {
    ModelRoot root(deepseek_vl2_root);
    EXPECT_EQ(SUCCESS, root.Init());
}

TEST(ModelRootTest, InitTest_NOT_FIND_FILE) {
    ModelRoot root(deepseek_vl2_root + "-no_exist");
    EXPECT_EQ(FILE_NOT_EXIST, root.Init());
}

TEST(ModelRootTest, CheckConfigResult_SUCCESS)
{
    ModelRoot root(deepseek_vl2_root);
    Ret ret = root.Init();
    EXPECT_EQ(SUCCESS, ret);
    Config conf = root.config();
    EXPECT_EQ(conf["global_view_pos"], "head");
    EXPECT_EQ(conf["language_config"]["architectures"][0],
            "DeepseekV2ForCausalLM");
    conf = root.processor_config();
    EXPECT_EQ(conf["downsample_ratio"], 2);
    EXPECT_EQ(conf["normalize"], true);

    conf = root.special_tokens_map();
    EXPECT_EQ(conf["pad_token"]["content"], "<｜▁pad▁｜>");

    conf = root.tokenizer_config();
    EXPECT_EQ(conf["add_bos_token"], true);
    EXPECT_EQ(conf["bos_token"], "<｜begin▁of▁sentence｜>");

    conf = root.tokenizer();
    EXPECT_EQ(conf["added_tokens"][0]["content"],
        "<｜begin▁of▁sentence｜>");
}

TEST(ModelRootTest, FindAllSafeTensorFiles_SUCCESS)
{
    ModelRoot root(deepseek_vl2_root);
    Ret ret = root.Init();
    EXPECT_EQ(SUCCESS, ret);

    ret = root.FindAllSafeTensorFiles();
    EXPECT_EQ(SUCCESS, ret);
}

TEST(ModelRootTest, FindAllSafeTensorFiles_FAILED)
{
    const std::string dir = "../3rd/googletest";
    ModelRoot root(dir);
    Ret ret = root.Init();
    EXPECT_EQ(FILE_NOT_EXIST, ret);

    ret = root.FindAllSafeTensorFiles();
    EXPECT_EQ(ret, SAFETENSOR_FILES_NOT_EXIST);
}